#!/usr/bin/env node

import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import {
  CallToolRequestSchema,
  ErrorCode,
  ListToolsRequestSchema,
  McpError,
} from "@modelcontextprotocol/sdk/types.js";
import { z } from "zod";

import { DocumentManager } from "./document-manager";
import { VectorSearchService } from "./vector-search";
import { StructuredSearchService } from "./structured-search";
import { BriefingService } from "./briefing-service";
import { AliBailianConfig } from "./ali-bailian-service";

// 初始化阿里云百炼配置
const aliBailianConfig: AliBailianConfig | undefined = (() => {
  const apiKey = process.env.ALI_BAILIAN_API_KEY;
  if (!apiKey) {
    console.error('Error: ALI_BAILIAN_API_KEY environment variable is required');
    console.error('Please set your Alibaba Cloud Bailian API key:');
    console.error('export ALI_BAILIAN_API_KEY="your-api-key-here"');
    process.exit(1);
  }
  
  return {
    apiKey,
    baseUrl: process.env.ALI_BAILIAN_BASE_URL,
  };
})();

// 初始化服务
const documentManager = new DocumentManager(aliBailianConfig);
const vectorSearchService = new VectorSearchService(aliBailianConfig);
const structuredSearchService = new StructuredSearchService();
const briefingService = new BriefingService(aliBailianConfig);

// 创建服务器实例
const server = new Server(
  {
    name: "cb-document-helper",
    version: "1.0.0",
  },
  {
    capabilities: {
      tools: {},
    },
  }
);

// 工具参数模式定义
const GetBriefingArgsSchema = z.object({
  filePath: z.string().describe("The absolute path to the markdown file"),
});

const SearchDocumentChunksArgsSchema = z.object({
  query: z.string().describe("Search query for document chunks"),
  topK: z.number().min(1).max(50).default(10).describe("Number of top results to return"),
});

const StructuredSearchArgsSchema = z.object({
  keywords: z.string().describe("Keywords to search in frontmatter metadata"),
  directoryPath: z.string().describe("Absolute path of root directory to search in"),
});

const RefreshDocumentsArgsSchema = z.object({
  directoryPath: z.string().describe("Absolute path of root directory of markdown files you want to update"),
  forceUpdate: z.boolean().default(false).describe("Force update all documents regardless of modification time"),
});

const StartRefreshTaskArgsSchema = z.object({
  directoryPath: z.string().describe("Absolute path of root directory of markdown files you want to update"),
  forceUpdate: z.boolean().default(false).describe("Force update all documents regardless of modification time"),
});

const GetTaskStatusArgsSchema = z.object({
  taskId: z.string().describe("Task ID returned from start_refresh_task"),
});

// 注册工具
server.setRequestHandler(ListToolsRequestSchema, async () => {
  return {
    tools: [
      {
        name: "get_document_briefing",
        description: "Get a briefing (summary) for a specific markdown document",
        inputSchema: {
          type: "object",
          properties: {
            filePath: {
              type: "string",
              description: "The path to the markdown file",
            },
          },
          required: ["filePath"],
        },
      },
      {
        name: "search_document_chunks",
        description: "Search for document chunks using semantic similarity",
        inputSchema: {
          type: "object",
          properties: {
            query: {
              type: "string",
              description: "Search query for document chunks",
            },
            topK: {
              type: "number",
              description: "Number of top results to return (1-50)",
              minimum: 1,
              maximum: 50,
              default: 10,
            },
          },
          required: ["query"],
        },
      },      {
        name: "structured_search",
        description: "Search documents based on frontmatter metadata (structured search)",
        inputSchema: {
          type: "object",
          properties: {
            keywords: {
              type: "string",
              description: "Keywords to search in frontmatter metadata",
            },
            directoryPath: {
              type: "string",
              description: "Absolute path of root directory to search in",
            },
          },
          required: ["keywords"],
        },
      },{
        name: "refresh_documents",
        description: "Refresh document chunks and briefings for all markdown files in a directory (synchronous, may timeout for large directories)",
        inputSchema: {
          type: "object",
          properties: {
            directoryPath: {
              type: "string",
              description: "Absolute path of root directory of markdown files you want to update",
            },
            forceUpdate: {
              type: "boolean",
              description: "Force update all documents regardless of modification time",
              default: false,
            },
          },
          required: [],
        },
      },
      {
        name: "start_refresh_task",
        description: "Start an asynchronous document refresh task (recommended for large directories)",
        inputSchema: {
          type: "object",
          properties: {
            directoryPath: {
              type: "string",
              description: "Absolute path of root directory of markdown files you want to update",
            },
            forceUpdate: {
              type: "boolean",
              description: "Force update all documents regardless of modification time",
              default: false,
            },
          },
          required: [],
        },
      },
      {
        name: "get_task_status",
        description: "Get the status of an asynchronous refresh task",
        inputSchema: {
          type: "object",
          properties: {
            taskId: {
              type: "string",
              description: "Task ID returned from start_refresh_task",
            },
          },
          required: ["taskId"],
        },
      },
      {
        name: "list_all_tasks",
        description: "List all refresh tasks and their current status",
        inputSchema: {
          type: "object",
          properties: {},
          required: [],
        },
      },
    ],
  };
});

server.setRequestHandler(CallToolRequestSchema, async (request) => {
  try {
    const { name, arguments: args } = request.params;

    switch (name) {
      case "get_document_briefing": {
        const { filePath } = GetBriefingArgsSchema.parse(args);
        const briefing = await briefingService.getBriefing(filePath);
        
        return {
          content: [
            {
              type: "text",
              text: briefing,
            },
          ],
        };
      }

      case "search_document_chunks": {
        const { query, topK } = SearchDocumentChunksArgsSchema.parse(args);
        const results = await vectorSearchService.searchChunks(query, topK);
        
        const formattedResults = results.map((result, index) => 
          `**Result ${index + 1}** (Score: ${result.score.toFixed(3)})\n` +
          `File: ${result.filePath}\n` +
          `Lines: ${result.startLine}-${result.endLine}\n` +
          `Content:\n${result.content}\n`
        ).join('\n---\n');

        return {
          content: [
            {
              type: "text",
              text: formattedResults || "No matching chunks found.",
            },
          ],
        };
      }      case "structured_search": {
        const { keywords, directoryPath } = StructuredSearchArgsSchema.parse(args);
        const results = await structuredSearchService.searchByMetadata(keywords, directoryPath || process.cwd());
        
        const formattedResults = results.map((result, index) => 
          `**Result ${index + 1}**\n` +
          `File: ${result.filePath}\n` +
          `Metadata: ${JSON.stringify(result.metadata, null, 2)}\n` +
          `Briefing: ${result.briefing}\n`
        ).join('\n---\n');

        return {
          content: [
            {
              type: "text",
              text: formattedResults || "No matching documents found.",
            },
          ],
        };
      }      case "refresh_documents": {
        const { directoryPath, forceUpdate } = RefreshDocumentsArgsSchema.parse(args);
        const stats = await documentManager.refreshDocuments(directoryPath || process.cwd(), forceUpdate);
        
        return {
          content: [
            {
              type: "text",
              text: `Document refresh completed:\n` +
                   `- Total files processed: ${stats.totalFiles}\n` +
                   `- Files updated: ${stats.updatedFiles}\n` +
                   `- Files skipped: ${stats.skippedFiles}\n` +
                   `- Chunks created: ${stats.chunksCreated}\n` +
                   `- Briefings generated: ${stats.briefingsGenerated}`,
            },
          ],
        };
      }

      case "start_refresh_task": {
        const { directoryPath, forceUpdate } = StartRefreshTaskArgsSchema.parse(args);
        const taskId = await documentManager.startRefreshDocumentsAsync(directoryPath || process.cwd(), forceUpdate);
        
        return {
          content: [
            {
              type: "text",
              text: `Refresh task started successfully.\nTask ID: ${taskId}\n\nUse get_task_status with this task ID to check progress.`,
            },
          ],
        };
      }

      case "get_task_status": {
        const { taskId } = GetTaskStatusArgsSchema.parse(args);
        const task = documentManager.getTaskStatus(taskId);
        
        if (!task) {
          return {
            content: [
              {
                type: "text",
                text: `Task ${taskId} not found.`,
              },
            ],
          };
        }

        let statusText = `Task ID: ${task.id}\n`;
        statusText += `Status: ${task.status}\n`;
        statusText += `Directory: ${task.directoryPath}\n`;
        statusText += `Force Update: ${task.forceUpdate}\n`;
        statusText += `Started: ${task.startTime.toISOString()}\n`;
        
        if (task.endTime) {
          statusText += `Ended: ${task.endTime.toISOString()}\n`;
          const duration = task.endTime.getTime() - task.startTime.getTime();
          statusText += `Duration: ${Math.round(duration / 1000)}s\n`;
        }

        statusText += `Progress: ${task.progress.processedFiles}/${task.progress.totalFiles} files\n`;
        
        if (task.progress.currentFile) {
          statusText += `Current file: ${task.progress.currentFile}\n`;
        }

        if (task.stats) {
          statusText += `\nResults:\n`;
          statusText += `- Total files processed: ${task.stats.totalFiles}\n`;
          statusText += `- Files updated: ${task.stats.updatedFiles}\n`;
          statusText += `- Files skipped: ${task.stats.skippedFiles}\n`;
          statusText += `- Chunks created: ${task.stats.chunksCreated}\n`;
          statusText += `- Briefings generated: ${task.stats.briefingsGenerated}\n`;
        }

        if (task.error) {
          statusText += `\nError: ${task.error}\n`;
        }

        return {
          content: [
            {
              type: "text",
              text: statusText,
            },
          ],
        };
      }

      case "list_all_tasks": {
        const tasks = documentManager.getAllTasks();
        
        if (tasks.length === 0) {
          return {
            content: [
              {
                type: "text",
                text: "No refresh tasks found.",
              },
            ],
          };
        }

        const taskList = tasks.map(task => {
          const duration = task.endTime 
            ? Math.round((task.endTime.getTime() - task.startTime.getTime()) / 1000)
            : Math.round((Date.now() - task.startTime.getTime()) / 1000);
          
          return `**${task.id}**\n` +
            `Status: ${task.status}\n` +
            `Directory: ${task.directoryPath}\n` +
            `Progress: ${task.progress.processedFiles}/${task.progress.totalFiles}\n` +
            `Duration: ${duration}s\n` +
            `Started: ${task.startTime.toLocaleString()}`;
        }).join('\n\n---\n\n');

        return {
          content: [
            {
              type: "text",
              text: `All Tasks:\n\n${taskList}`,
            },
          ],
        };
      }

      default:
        throw new McpError(
          ErrorCode.MethodNotFound,
          `Unknown tool: ${name}`
        );
    }
  } catch (error) {
    if (error instanceof z.ZodError) {
      throw new McpError(
        ErrorCode.InvalidParams,
        `Invalid parameters: ${error.errors.map(e => `${e.path.join('.')}: ${e.message}`).join(', ')}`
      );
    }
    
    console.error("Tool execution error:", error);
    throw new McpError(
      ErrorCode.InternalError,
      error instanceof Error ? error.message : "Unknown error occurred"
    );
  }
});

// 启动服务器
async function main() {
  const transport = new StdioServerTransport();
  await server.connect(transport);
  console.error("CB Document Helper MCP Server running on stdio");
}

if (require.main === module) {
  main().catch((error) => {
    console.error("Fatal error:", error);
    process.exit(1);
  });
}
